#include "meshTypes.hpp"
#include "fieldFunctions.hpp"
#include "functionParserWrapper.hpp"
#include "field.hpp"
#include "simulationTime.hpp"
#include <iostream>

namespace hms
{

auto findValuesAtTime( const TimeSeries& values, scalar time )
	-> TimeSeries::const_iterator
{
	auto valsIter = std::upper_bound(
		values.cbegin(), values.cend(), 
		std::pair<scalar, ArrayXXs>{ time, {} },
		[](const auto& a, const auto& b){ return a.first < b.first; }
	) - 1;
	if ( valsIter < values.cbegin() ){
		std::cerr <<
			"WARNING: No values available for time t = " << time << "s! "
			"Values for minimum time index t = " << values.cbegin()->first << 
			"s are used instead.\n";
		return values.cbegin();
	} else
		return valsIter;
}

template<typename MeshType>
void resizeFieldValues(
	ArrayXXs& values,
	Index nComponents,
	const MeshBase<MeshType>& mesh,
	Uniformity isUniform
){
	if (isUniform)
		resizeFieldValues<Uniform>(values, nComponents, mesh);
	else
		resizeFieldValues<NonUniform>(values, nComponents, mesh);
}

template<typename MeshType>
void setCorners(
	[[maybe_unused]] ArrayXXs& values,
	[[maybe_unused]] const MeshBase<MeshType>& mesh,
	[[maybe_unused]] scalar cornerValue
){
	if constexpr ( std::is_base_of_v<StructMeshBase<MeshType>, MeshType> ){
		static constexpr Index ng {1};
		Index
			nx { mesh.derived().nCellsX() },
			ny { mesh.derived().nCellsY() },
			nc { values.rows() };

		auto reshapedVals { values.reshaped( (nx + 2*ng)*nc, ny + 2*ng ) };
		reshapedVals.bottomLeftCorner ( ng*nc, ng ) = cornerValue;
		reshapedVals.bottomRightCorner( ng*nc, ng ) = cornerValue;
		reshapedVals.topLeftCorner    ( ng*nc, ng ) = cornerValue;
		reshapedVals.topRightCorner   ( ng*nc, ng ) = cornerValue;
	}
}

namespace internal {

template<Uniformity isUniform, bool useTime, bool useZField, typename MeshType = UniMesh>
void setFieldByFunctions(
	ArrayXXs& vals,
	const std::vector<FunctionParserWrapper>& functions,
	[[maybe_unused]] std::experimental::observer_ptr<const MeshBase<MeshType>> mesh,
	[[maybe_unused]] std::experimental::observer_ptr<const Time> time,
	[[maybe_unused]] std::experimental::observer_ptr<const Field<MeshType>> zField
){
	assert(
		functions.size() ==
		static_cast<decltype(functions.size())>( vals.rows() )
	);
	scalar t {0}, dt {0};
	if constexpr (useTime){
		assert( time );
		t = time->current();
		dt = time->step();
	}
	if constexpr (useZField)
		assert( zField );
	
	if constexpr (isUniform){
		for ( Index i{0}; i<vals.rows(); ++i ){
			vals(i,0) = functions[i].eval(
				0,0,0,t,dt
			);
		}
	} else {
		assert( mesh );
		mesh->forEachDomainCell( [&](const auto& cell){
			ArrayXXs::ColXpr col { vals.col( cell.index() ) };
			const auto& centroid { cell.centroid() };

			for ( size_t i{0}; i<functions.size(); ++i ){
				col(i) =
				functions[i].eval(
					centroid.x(), centroid.y(),
					useZField ? zField->values()( 0,cell.index() ) : 0,
					t, dt
				);
			}
		});
	}
}

}

template<Uniformity isUniform, typename MeshType>
void setFieldByFunctions(
	ArrayXXs& vals,
	const std::vector<FunctionParserWrapper>& functions,
	const MeshBase<MeshType>& mesh,
	const Time& time,
	std::experimental::observer_ptr<const Field<MeshType>> zField
){
	internal::setFieldByFunctions<isUniform, true, !isUniform>(
		vals, functions,
		std::experimental::make_observer( &mesh ),
		std::experimental::make_observer( &time ),
		zField
	);
}

template<Uniformity isUniform, typename MeshType>
void setFieldByFunctions(
	ArrayXXs& vals,
	const std::vector<FunctionParserWrapper>& functions,
	const MeshBase<MeshType>& mesh,
	const Field<MeshType>& z,
	bool useZField
){
	if ( useZField )
		internal::setFieldByFunctions<isUniform, false, true>(
			vals, functions,
			std::experimental::make_observer( &mesh ),
			nullptr,
			std::experimental::make_observer( &z )
		);
	else
		internal::setFieldByFunctions<isUniform, false, false, MeshType>(
			vals, functions,
			std::experimental::make_observer( &mesh ),
			nullptr,
			nullptr
		);
}

#define INSTANTIATE_UNIFORMITY(IS_UNIFORM, MeshType) \
	template void setFieldByFunctions<IS_UNIFORM, MeshType>( \
		ArrayXXs&, const std::vector<FunctionParserWrapper>&, \
		const MeshBase<MeshType>&, const Time& time, \
		std::experimental::observer_ptr<const Field<MeshType>> ); \
	template void setFieldByFunctions<IS_UNIFORM, MeshType>( \
		ArrayXXs& vals, const std::vector<FunctionParserWrapper>&, \
		const MeshBase<MeshType>&, const Field<MeshType>&, bool );

#define INSTANTIATE_FOR_ALL_MESHTYPES(MeshType) \
	INSTANTIATE_UNIFORMITY(Uniform, MeshType) \
	INSTANTIATE_UNIFORMITY(NonUniform, MeshType) \
	template void resizeFieldValues( \
		ArrayXXs&, Index, const MeshBase<MeshType>&, Uniformity); \
	template void setCorners( ArrayXXs&, const MeshBase<MeshType>&, scalar );

#include "meshInstantiator.hpp"
#undef INSTANTIATE_UNIFORMITY

} // end namespace hms
